import torch
from ultralytics import YOLO
import cv2
import numpy as np

# 设备
device = torch.device('cpu')
print(f"Using device: {device}")

# 加载模型
model1_path = 'data/weights/best.torchscript'
model2_path = 'data/weights/yolopv2.pt'
model1 = torch.jit.load(model1_path).to(device)
model2 = torch.jit.load(model2_path).to(device)

# 加载视频
video_path = "C:/Users/永生理想/Desktop/4.mp4"
cap = cv2.VideoCapture(video_path)

width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)

# 设置视频写入器
output_path = "C:/Users/永生理想/Desktop/output_video.mp4"
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

cv2.namedWindow('YOLOv8 Detection', cv2.WINDOW_NORMAL)

# 类别名称（需手动提供）
class_names1 = ['class1', 'class2']  # 替换为 best.pt 的类别
class_names2 = ['class1', 'class2']  # 替换为 yolopv2.pt 的类别

# 模型期望的输入尺寸
input_size = 640

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # 预处理帧
    img = cv2.resize(frame, (input_size, input_size), interpolation=cv2.INTER_LINEAR)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
    img = img.unsqueeze(0)

    # 第一个模型推理
    with torch.no_grad():
        preds1 = model1(img)
        print("preds1 shape:", preds1.shape)  # 调试输出形状
    # 假设输出为 [1, N, 6] 格式：[x1, y1, x2, y2, conf, cls]
    for pred in preds1[0]:  # preds1[0] 是 [N, 6]
        x1, y1, x2, y2, conf, cls = pred
        if conf > 0.5:
            x1, x2 = int(x1 * width / input_size), int(x2 * width / input_size)
            y1, y2 = int(y1 * height / input_size), int(y2 * height / input_size)
            cls = int(cls)
            label = f'{class_names1[cls]} {conf:.2f}'
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    # 第二个模型推理
    with torch.no_grad():
        preds2 = model2(img)
        print("preds2 shape:", preds2.shape)  # 调试输出形状
    for pred in preds2[0]:  # preds2[0] 是 [N, 6]
        x1, y1, x2, y2, conf, cls = pred
        if conf > 0.5:
            x1, x2 = int(x1 * width / input_size), int(x2 * width / input_size)
            y1, y2 = int(y1 * height / input_size), int(y2 * height / input_size)
            cls = int(cls)
            label = f'{class_names2[cls]} {conf:.2f}'
            cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

    cv2.imshow('YOLOv8 Detection', frame)
    out.write(frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
out.release()
cv2.destroyAllWindows()